#pragma once

#include <iostream>
#include <cassert>

using namespace std;

template <class K>
struct AVLNode
{
	K _data;
	AVLNode* _left;
	AVLNode* _right;
	AVLNode* _parent;
	int _bf;

	AVLNode(K data)
		:_data(data),
		 _left(nullptr),
		 _right(nullptr),
		 _parent(nullptr),
		 _bf(0)
	{}
};

template <class K>
class AVLTree
{
public:
	using Node = AVLNode<K>;

	bool Insert(K key)
	{
		Node* newnode = new Node(key);
		if (_root == nullptr) _root = newnode;
		else
		{
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				parent = cur;
				if (cur->_data == key) return false;
				else if (cur->_data < key) cur = cur->_right;
				else cur = cur->_left;
			}
			newnode->_parent = parent;
			if (parent->_data > key)
			{
				parent->_left = newnode;
				parent->_bf -= 1;
			}
			else
			{
				parent->_right = newnode;
				parent->_bf += 1;
			}
			while (parent)
			{
				if (parent->_bf == 0)
				{
					break;
				}
				else if (parent->_bf == -1 || parent->_bf == 1)
				{
					if (parent->_parent)
					{
						if (parent == parent->_parent->_left)
							parent->_parent->_bf -= 1;
						else
							parent->_parent->_bf += 1;
					}
					parent = parent->_parent;
				}
				else if (parent->_bf == -2 || parent->_bf == 2)
				{
					if (parent->_bf == -2 && parent->_left->_bf == -1)
					{
						RotateR(parent);
					}
					else if (parent->_bf == 2 && parent->_right->_bf == 1)
					{
						RotateL(parent);
					}
					else if (parent->_bf == -2 && parent->_left->_bf == 1)
					{
						RotateLR(parent);
					}
					else if (parent->_bf == 2 && parent->_right->_bf == -1)
					{
						RotateRL(parent);
					}
					else
					{
						assert(false);
					}
					break;
				}
				else
				{
					assert(false);
				}
			}
		}
		return true;
	}

	Node* Find(K key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (cur->_data == key) return cur;
			else if (cur->_data < key) cur = cur->_right;
			else cur = cur->_left;
		}
		return nullptr;
	}

	bool IsBalanceTree()
	{
		return _IsBalanceTree(_root);
	}

	void Inorder()
	{
		_Inorder(_root);
		cout << endl;
	}

private:

	void _Inorder(Node* root)
	{
		if (root == nullptr) return;
		_Inorder(root->_left);
		cout << root->_data << " ";
		_Inorder(root->_right);
	}

	int _Height(Node* root)
	{
		if (root == nullptr)
			return 0;

		int heightl = _Height(root->_left);
		int heightr = _Height(root->_right);
		return  heightl > heightr ? heightl + 1 : heightr + 1;
	}

	bool _IsBalanceTree(Node* root)
	{
		if (root == nullptr)
			return true;

		int heightL = _Height(root->_left);
		int heightR = _Height(root->_right);
		int diff = heightR - heightL;

		if (diff != root->_bf)
		{
			return false;
		}
		else if (abs(diff) >= 2)
		{
			return false;
		}
		else
		{
			return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
		}
	}

private:
	void RotateR(Node* parent)
	{
		Node* subl = parent->_left;
		Node* sublr = subl->_right;
		Node* pparent = parent->_parent;

		if (parent == _root) _root = subl;
		else if (pparent->_left == parent)
			pparent->_left = subl;
		else
			pparent->_right = subl;

		subl->_parent = pparent;
		subl->_right = parent;
		parent->_parent = subl;
		parent->_left = sublr;

		subl->_bf = 0;
		parent->_bf = 0;
	}

	void RotateL(Node* parent)
	{
		Node* subr = parent->_right;
		Node* subrl = subr->_left;
		Node* pparent = parent->_parent;
		
		if (parent == _root) _root = subr;
		else if (pparent->_left == parent)
			pparent->_left = subr;
		else
			pparent->_right = subr;

		subr->_parent = pparent;
		subr->_left = parent;
		parent->_parent = subr;
		parent->_right = subrl;

		subr->_bf = 0;
		parent->_bf = 0;
	}

	void RotateLR(Node* parent)
	{
		Node* subl = parent->_left;
		Node* sublr = subl->_right;
		int bf = sublr->_bf;

		RotateL(subl);
		RotateR(parent);

		sublr->_bf = 0;
		if(bf == 1)
		{ 
			subl->_bf = -1;
			parent->_bf = 0;
		}
		else if (bf == -1)
		{
			subl->_bf = 0;
			parent->_bf = 1;
		}
		else
		{
			subl->_bf = 0;
			parent->_bf = 0;
		}
	}

	void RotateRL(Node* parent)
	{
		Node* subr = parent->_right;
		Node* subrl = subr->_left;
		int bf = subrl->_bf;

		RotateR(subr);
		RotateL(parent);

		subrl->_bf = 0;
		if (bf == 1)
		{
			subr->_bf = 0;
			parent->_bf = -1;
		}
		else if (bf == -1)
		{
			subr->_bf = 1;
			parent->_bf = 0;
		}
		else
		{
			subr->_bf = 0;
			parent->_bf = 0;
		}
	}
private:
	Node* _root = nullptr;
};